import os
import gluoncv
os.environ["MXNET_CUDNN_AUTOTUNE_DEFAULT"] = "0"
import mxnet as mx
from lib.nms.nms import gpu_nms_wrapper
from lib.config import get_default_config
import os,cv2,time
import numpy as np
from models.ml_resnet101_ori import RFCN_Resnet
import mxnet.ndarray as nd
from lib.common import lsdir
import matplotlib.pyplot as plt
from lib.data.coco import COCODetection
from lib.common import shortname
from tqdm import tqdm
def main():
    device_id = 4
    nms_wrapper = gpu_nms_wrapper(thresh=0.3,device_id=device_id)
    cfg = get_default_config(network_stride = [16,16],scales=[4,8,16,32],ratios=(.125,0.25,0.5,1))
    cfg.dataset.NUM_CLASSES = 3#with background
    # cfg.pretrained = "/data1/zyx/yks/object_detection/Deformable-ConvNets/output/rfcn_dcn/coco/resnet_v1_101_coco_trainval_rfcn_dcn_end2end_ohem/formual_train/rfcn_dcn_coco-0007.params" #load resnet-101 pretrained model.
    cfg.resume = "output/no_latex/faster_rcnn_resnet50_v1b_custom_0017_0.8273.params"
    cfg.TEST.KEEP_THRESHOLD = .7
    cfg.TRAIN.BBOX_NORMALIZATION_PRECOMPUTED = False if cfg.resume is None else True

    net = RFCN_Resnet(cfg = cfg,viz_env = None)
    # net.initialize()
    net.collect_params().reset_ctx(ctx=mx.gpu(device_id))

    # val_anno_path = "../../dataset/coco/annotations/instances_val2014.json"
    # val_image_root = "../../dataset/coco/images/val2014"
    # val_dataset = COCODetection(anno_path=val_anno_path,image_root=val_image_root)
    # images2val = [val_dataset.at_with_image_path(i)[0] for i in range(1000) ]
    detection_result = []
    images2val = list(lsdir("/data1/zyx/yks/dataset/ocr_formula/val"))
    for img_name in tqdm(images2val):
        t0 = time.time()
        ori_img = cv2.imread(img_name)[:,:,:]
        fscale = 1024.0 / ori_img.shape[1]
        img_resized = cv2.resize(ori_img, (0, 0), fx=fscale, fy=fscale)

        img_float = img_resized.astype(np.float32)
        mean = np.array([103.06,115.90, 123.15])[np.newaxis,np.newaxis]
        img_float -= mean
        img_float = img_float[:,:,(2,1,0)]
        img_float = np.transpose(img_float,(2,0,1))
        data = mx.nd.array(img_float[np.newaxis],ctx = mx.gpu(device_id))

        im_info = nd.array([[data.shape[2],data.shape[3],3]],ctx=mx.gpu(device_id))
        rois, cls_score, bbox_pred = net(img_resized[np.newaxis],data,im_info,None,None,None,None,None)
        pred_bboxes,pred_scores,pred_labels = net.post_process(rois, bbox_pred, cls_score,nms_wrapper,data.shape,scale=fscale)
        # print(time.time()-t0)
        gluoncv.utils.viz.plot_bbox(ori_img[:,:,::-1],bboxes = pred_bboxes,scores = pred_scores ,class_names=["formual","picture","handwriting"],labels=pred_labels)
        plt.show()
    #     for bbox,cls,label in zip(pred_bboxes,pred_scores,pred_labels):
    #         one_result = [shortname(img_name)]+list(bbox)+[label,cls]
    #         detection_result.append(one_result)
    # result = val_dataset.coco_detection_evaluate(detection_result)
    pass
if __name__ == '__main__':
    main()
